function [capacity_mimo, capacity_scaling, fig_mimo] = mimo_capacity(EbN0_dB, num_antennas, channel_types)
% MIMO信道容量分析
% 输入参数：
%   EbN0_dB - 信噪比范围 (dB)
%   num_antennas - 天线数量数组 [Nt, Nr]
%   channel_types - 信道类型 ('iid', 'correlated', 'rice')
% 输出参数：
%   capacity_mimo - MIMO容量
%   capacity_scaling - 容量缩放分析
%   fig_mimo - 图形句柄

% 添加路径
addpath('../Common');
colors = color_definitions();

% 参数设置
bandwidth = 180e3; % 带宽 (Hz)
noise_power_dBm = -174 + 10*log10(bandwidth);
noise_power = 10^((noise_power_dBm-30)/10);
num_realizations = 100; % 信道实现次数

% 初始化结果数组
num_snr = length(EbN0_dB);
num_antenna_cases = size(num_antennas, 1);
num_channel_types = length(channel_types);
capacity_mimo = zeros(num_snr, num_antenna_cases, num_channel_types);
capacity_scaling = zeros(num_antenna_cases, 2); % [理论缩放, 实际缩放]

fprintf('MIMO信道容量分析...\n');

% MIMO容量计算
for snr_idx = 1:num_snr
    snr_linear = 10^(EbN0_dB(snr_idx)/10);
    
    for ant_idx = 1:num_antenna_cases
        Nt = num_antennas(ant_idx, 1); % 发射天线数
        Nr = num_antennas(ant_idx, 2); % 接收天线数
        
        for ch_idx = 1:num_channel_types
            channel_type = channel_types{ch_idx};
            
            % 生成MIMO信道矩阵
            H = zeros(Nr, Nt, num_realizations);
            
            for realization = 1:num_realizations
                switch lower(channel_type)
                    case 'iid'
                        % IID瑞利衰落信道
                        H(:, :, realization) = sqrt(0.5) * (randn(Nr, Nt) + 1i * randn(Nr, Nt));
                        
                    case 'correlated'
                        % 相关信道 (发射和接收相关)
                        % 发射相关矩阵
                        Rt = toeplitz([1, 0.5, zeros(1, Nt-2)]);
                        % 接收相关矩阵  
                        Rr = toeplitz([1, 0.3, zeros(1, Nr-2)]);
                        
                        % 生成相关信道
                        H_uncorrelated = sqrt(0.5) * (randn(Nr, Nt) + 1i * randn(Nr, Nt));
                        H(:, :, realization) = Rr^(1/2) * H_uncorrelated * Rt^(1/2);
                        
                    case 'rice'
                        % 莱斯信道 (K因子=3dB)
                        K_dB = 3;
                        K_linear = 10^(K_dB/10);
                        LOS_component = sqrt(K_linear/(K_linear+1));
                        NLOS_component = sqrt(1/(K_linear+1));
                        
                        % LOS分量 (假设为全1矩阵)
                        H_LOS = ones(Nr, Nt) / sqrt(Nr*Nt);
                        H_NLOS = sqrt(0.5) * (randn(Nr, Nt) + 1i * randn(Nr, Nt));
                        
                        H(:, :, realization) = LOS_component * H_LOS + NLOS_component * H_NLOS;
                        
                    otherwise
                        % 默认为IID信道
                        H(:, :, realization) = sqrt(0.5) * (randn(Nr, Nt) + 1i * randn(Nr, Nt));
                end
            end
            
            % 计算MIMO容量 (使用奇异值分解)
            total_capacity = 0;
            for realization = 1:num_realizations
                H_real = H(:, :, realization);
                
                % 奇异值分解
                [U, S, V] = svd(H_real);
                singular_values = diag(S);
                
                % 计算容量 (注水算法简化版)
                capacity_realization = 0;
                for stream = 1:min(Nt, Nr)
                    if singular_values(stream) > 0
                        % 等功率分配
                        power_per_stream = snr_linear / Nt;
                        stream_snr = power_per_stream * singular_values(stream)^2 / noise_power;
                        
                        if stream_snr > 0
                            capacity_realization = capacity_realization + bandwidth * log2(1 + stream_snr);
                        end
                    end
                end
                
                total_capacity = total_capacity + capacity_realization;
            end
            
            % 平均容量
            capacity_mimo(snr_idx, ant_idx, ch_idx) = total_capacity / num_realizations;
        end
    end
end

% 容量缩放分析
for ant_idx = 1:num_antenna_cases
    Nt = num_antennas(ant_idx, 1);
    Nr = num_antennas(ant_idx, 2);
    
    % 理论缩放 (min(Nt, Nr))
    capacity_scaling(ant_idx, 1) = min(Nt, Nr);
    
    % 实际缩放 (相对于单天线)
    snr_mid = round(num_snr/2);
    capacity_scaling(ant_idx, 2) = capacity_mimo(snr_mid, ant_idx, 1) / capacity_mimo(snr_mid, 1, 1);
end

%% 可视化结果
fig_mimo = figure('Name', 'MIMO信道容量分析', 'Position', [100, 100, 1200, 800]);

% 子图1: 容量 vs SNR (不同天线配置)
subplot(2, 2, 1);
for ant_idx = 1:num_antenna_cases
    Nt = num_antennas(ant_idx, 1);
    Nr = num_antennas(ant_idx, 2);
    plot(EbN0_dB, capacity_mimo(:, ant_idx, 1)/1e6, 'LineWidth', 2, 'Color', colors(ant_idx, :));
    hold on;
end
grid on;
xlabel('Eb/N0 (dB)');
ylabel('容量 (Mbps)');
title('MIMO容量 vs 信噪比 (IID信道)');
legend(arrayfun(@(x) sprintf('%dx%d', num_antennas(x,1), num_antennas(x,2)), 1:num_antenna_cases, 'UniformOutput', false), 'Location', 'NorthWest');

% 子图2: 容量 vs 天线数量 (固定SNR)
subplot(2, 2, 2);
snr_fixed = round(num_snr/2);
antenna_numbers = min(num_antennas, [], 2); % 取发射和接收的最小值
capacities_fixed_snr = squeeze(capacity_mimo(snr_fixed, :, 1));
plot(antenna_numbers, capacities_fixed_snr/1e6, 'o-', 'LineWidth', 2, 'Color', colors(1, :));
hold on;

% 理论线性缩放
linear_scaling = antenna_numbers * capacities_fixed_snr(1) / antenna_numbers(1);
plot(antenna_numbers, linear_scaling/1e6, 'k--', 'LineWidth', 2);
grid on;
xlabel('天线数量 (min(Nt, Nr))');
ylabel('容量 (Mbps)');
title(sprintf('MIMO容量 vs 天线数量 (Eb/N0 = %d dB)', EbN0_dB(snr_fixed)));
legend({'仿真结果', '线性缩放'}, 'Location', 'NorthWest');

% 子图3: 不同信道类型的影响
subplot(2, 2, 3);
ant_fixed = 3; % 4x4 MIMO
for ch_idx = 1:num_channel_types
    plot(EbN0_dB, capacity_mimo(:, ant_fixed, ch_idx)/1e6, 'LineWidth', 2, 'Color', colors(ch_idx, :));
    hold on;
end
grid on;
xlabel('Eb/N0 (dB)');
ylabel('容量 (Mbps)');
title(sprintf('信道类型对MIMO容量的影响 (%dx%d)', ...
              num_antennas(ant_fixed,1), num_antennas(ant_fixed,2)));
legend(channel_types, 'Location', 'NorthWest');

% 子图4: 容量缩放分析
subplot(2, 2, 4);
bar(antenna_numbers, [capacity_scaling(:, 1), capacity_scaling(:, 2)], 'FaceColor', colors(1, :));
legend({'理论缩放', '实际缩放'}, 'Location', 'NorthWest');
xlabel('天线配置 (min(Nt, Nr))');
ylabel('容量缩放因子');
title('MIMO容量缩放分析');

fprintf('MIMO容量分析完成！\n');

end